#include <iostream>
#include <fstream>
#include <assert.h>
#include "KMeans.h"
#include "CHMM.h"
#include "utils.hpp"
using namespace std;

int main()
{
    int seq_num = 4; //Number of samples
    const int state_num = 3; //State number
    const int mix_num = 1; //GMM mixture number
    const int dim_num = 2; //GMM dimension of feature

    double sample1[dim_num] = {0.5, 2.0};
    double sample2[dim_num] = {-1.0, 2.0};
    double sample3[dim_num] = {10.5, -2.0};
    double sample4[dim_num] = {8.5, 1.0};
    std::vector<double*> test_data_seq = {sample1, sample2, sample3, sample4}; //samples (seq_num)

    double start_prob[state_num] = {0.6, 0.3, 0.1};
    double transmat[state_num][state_num] = {0.7, 0.3, 0.0,
                                             0.3, 0.5, 0.2,
                                             0.2, 0.3, 0.5};
    double prior[mix_num] = {1.0};
    std::vector<double*> priors = {prior, prior , prior}; //priors (state_num)

    double means1[mix_num][dim_num] = {0.0, 0.0};
    double means2[mix_num][dim_num] = {-1.0, 5.0};
    double means3[mix_num][dim_num] = {9.0, 10.0};
    double means4[mix_num][dim_num] = {11.0, -1.0};
    std::vector<double**> means = {(double**)means1, (double**)means2, (double**)means3}; //means (state_num)

    double vars1[mix_num][dim_num] = {0.5, 0.5};
    std::vector<double**> covars = {(double**)vars1, (double**)vars1, (double**)vars1}; //vars (state_num)

    //Test CHMM
    CHMM* chmm = new CHMM(state_num, dim_num, mix_num);

    chmm->SetStateInit(start_prob);
    chmm->SetStateTrans((double**)transmat, state_num);
    chmm->SetStateModel(priors, means, covars, mix_num, dim_num);
    //Load parameters from config profiles
//    ifstream chmm_file_in("chmm.conf");
//    assert(chmm_file_in);
//    chmm_file_in >> *chmm;
//    chmm_file_in.close();

    vector<int> state;
    double prob = chmm->Decode(test_data_seq, state);
    printf("The Probability is %e\n", prob);
    printf("state: ");
    for(int i = 0; i < state.size(); ++i)
        printf("%d ", state[i]);
    printf("\n");

    seq_num = test_data_seq.size();
    double** prob_post = Alloc_Matrix(seq_num, state_num, 0.0);
    double prob_predict[state_num];
    chmm->Alloc_TempProbMatrix(seq_num);

    chmm->GetProbPosterior(test_data_seq, prob_post, prob_predict);
    Print_Matrix(prob_post, seq_num, state_num);
    printf("prob_predict:\n");
    Print_Vector(prob_predict, state_num);

    chmm->Free_TempProbMatrix(seq_num);
    Free_Matrix(prob_post, seq_num);


    //save GMM to file
    ofstream chmm_file_out("chmm.conf");
    assert(chmm_file_out);
    chmm_file_out << *chmm;
    chmm_file_out.close();

    delete chmm;

    return EXIT_SUCCESS;
}
